import torch
import math
from typing import Tuple


def hat(v: torch.Tensor) -> torch.Tensor:
    """
    将向量转换为斜对称矩阵 (hat operator)

    Args:
        v: 形状为 (..., 3) 的张量
    Returns:
        形状为 (..., 3, 3) 的斜对称矩阵
    """
    if v.shape[-1] != 3:
        raise ValueError("Input tensor last dimension must be 3")

    # 初始化斜对称矩阵
    shape = list(v.shape[:-1]) + [3, 3]
    skew = torch.zeros(shape, device=v.device, dtype=v.dtype)

    # 填充斜对称元素
    skew[..., 0, 1] = -v[..., 2]
    skew[..., 0, 2] = v[..., 1]
    skew[..., 1, 0] = v[..., 2]
    skew[..., 1, 2] = -v[..., 0]
    skew[..., 2, 0] = -v[..., 1]
    skew[..., 2, 1] = v[..., 0]

    return skew


def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    将轴角表示转换为旋转矩阵（使用罗德里格斯公式）

    Args:
        axis_angle: 形状为 (..., 3) 的张量，轴角表示
    Returns:
        形状为 (..., 3, 3) 的旋转矩阵
    """
    if axis_angle.shape[-1] != 3:
        raise ValueError("Input tensor last dimension must be 3")

    # 计算旋转角度
    angle = torch.norm(axis_angle, dim=-1, keepdim=True)

    # 处理零角度的情况
    mask = angle.squeeze(-1) < 1e-8
    if mask.any():
        # 对于零角度，返回单位矩阵
        shape = list(axis_angle.shape[:-1]) + [3, 3]
        identity = torch.eye(3, device=axis_angle.device, dtype=axis_angle.dtype)
        identity = identity.expand(shape)

        # 对于非零角度，计算旋转矩阵
        if not mask.all():
            non_zero_axis_angle = axis_angle[~mask]
            non_zero_angle = angle[~mask]

            # 计算单位旋转轴
            axis = non_zero_axis_angle / non_zero_angle

            # 计算斜对称矩阵
            skew = hat(axis)

            # 罗德里格斯公式
            cos_angle = torch.cos(non_zero_angle)
            sin_angle = torch.sin(non_zero_angle)

            non_zero_R = (torch.eye(3, device=axis_angle.device, dtype=axis_angle.dtype) +
                          sin_angle.unsqueeze(-1) * skew +
                          (1 - cos_angle).unsqueeze(-1) * torch.matmul(skew, skew))

            result = identity.clone()
            result[~mask] = non_zero_R
            return result
        else:
            return identity
    else:
        # 所有角度都非零
        axis = axis_angle / angle
        skew = hat(axis)

        cos_angle = torch.cos(angle)
        sin_angle = torch.sin(angle)

        R = (torch.eye(3, device=axis_angle.device, dtype=axis_angle.dtype) +
             sin_angle.unsqueeze(-1) * skew +
             (1 - cos_angle).unsqueeze(-1) * torch.matmul(skew, skew))

        return R


def _so3_exp_map(
        log_rot: torch.Tensor, eps: float = 0.0001
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    A helper function that computes the so3 exponential map and,
    apart from the rotation matrix, also returns intermediate variables
    that can be re-used in other functions.

    Args:
        log_rot: Batch of vectors of shape `(minibatch, 3)`.
        eps: A float constant handling the conversion singularity.

    Returns:
        Tuple of (R, rot_angles, skews, skews_square)
    """
    if len(log_rot.shape) != 2 or log_rot.shape[1] != 3:
        raise ValueError("Input tensor shape has to be Nx3.")

    # 计算旋转角度
    nrms = (log_rot * log_rot).sum(1)
    rot_angles = torch.clamp(nrms, eps).sqrt()

    # 计算斜对称矩阵及其平方
    skews = hat(log_rot)
    skews_square = torch.bmm(skews, skews)

    # 使用轴角转换得到旋转矩阵
    R = axis_angle_to_matrix(log_rot)

    return R, rot_angles, skews, skews_square


def so3_exp_map(log_rot: torch.Tensor, eps: float = 0.0001) -> torch.Tensor:
    """
    Convert a batch of logarithmic representations of rotation matrices `log_rot`
    to a batch of 3x3 rotation matrices using Rodrigues formula [1].

    In the logarithmic representation, each rotation matrix is represented as
    a 3-dimensional vector (`log_rot`) who's l2-norm and direction correspond
    to the magnitude of the rotation angle and the axis of rotation respectively.

    The conversion has a singularity around `log(R) = 0`
    which is handled by clamping controlled with the `eps` argument.

    Args:
        log_rot: Batch of vectors of shape `(minibatch, 3)`.
        eps: A float constant handling the conversion singularity.

    Returns:
        Batch of rotation matrices of shape `(minibatch, 3, 3)`.

    Raises:
        ValueError if `log_rot` is of incorrect shape.

    [1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
    """
    return _so3_exp_map(log_rot, eps=eps)[0]


# 测试代码
if __name__ == "__main__":
    # 测试示例
    batch_size = 2
    log_rot = torch.tensor([[0.1, 0.2, 0.3], [0.0, 0.0, 0.0]], dtype=torch.float32)

    # 使用替代实现
    R = so3_exp_map(log_rot)
    print("Rotation matrices:")
    print(R)

    # 测试辅助函数
    R_full, angles, skews, skews_sq = _so3_exp_map(log_rot)
    print("\nRotation angles:", angles)
    print("\nSkew-symmetric matrices:")
    print(skews)